#!/usr/bin/env python3

"""
DuoMAG MEP analyzer

Usage:
    $ ./duomag.py mep_file.ascii
"""

import sys
import os
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import scipy
import pandas as pd

program_version = "1.0"
print(f"DuoMAG MEP analyzer v{program_version}")
print("========================\n")

print("Imported modules:")
print(" Matplotlib " + matplotlib.__version__)
print("      NumPy " + np.__version__)
print("     pandas " + pd.__version__)
print("      SciPy " + scipy.__version__ + "\n")

if (len(sys.argv) == 0):
    print("ERROR: No input file provided!")
    quit()

file_name = sys.argv[1]

if not os.path.exists(file_name):
    print(f"ERROR: File {file_name} does not exist!")
    quit()

print(f"Reading file: {file_name}")
with open(file_name, encoding="utf8", errors='ignore') as f:
    ignore = f.readline()
    sw_name = f.readline().strip().split("=")[1]
    sw_version = f.readline().strip().split("=")[1]
    subject = f.readline().strip().split("=")[1]
    subject_id = f.readline().strip().split("=")[1]
    record_id = f.readline().strip().split("=")[1]
    record_created = f.readline().strip().split("=")[1]
    method = f.readline().strip().split("=")[1]
    ignore = f.readline()
    marker_latency_unit = f.readline().strip().split("=")[1]
    ignore = f.readline()
    sampling_interval = int(f.readline().strip().split("=")[1])
    sampling_interval_unit = f.readline().strip().split("=")[1]
    resolution = float(f.readline().strip().split("=")[1].replace(",", "."))
    resolution_unit = f.readline().strip().split("=")[1]
    signal_count = int(f.readline().strip().split("=")[1])
    stimulation = list(map(int, f.readline().strip().split(" ")[1:]))
    channel = list(map(int, f.readline().strip().split(" ")[1:]))
    samples_count = list(map(int, f.readline().strip().split(" ")[1:]))
    stim_sample = list(map(int, f.readline().strip().split(" ")[1:]))
    stim_intens = list(map(int, f.readline().strip().split(" ")[1:]))
    coil_type = f.readline().strip().split(" ")[1:]
    ignore = f.readline()
    markers_neg = np.array(f.readline().strip().replace(",", ".").split(" ")[1:])
    markers_pos = np.array(f.readline().strip().replace(",", ".").split(" ")[1:])
    ignore = f.readline()
    # data matrix: signals × samples
    mep_signal = np.zeros((samples_count[0], signal_count))
    for idx in range(mep_signal.shape[0]):
        mep_signal[idx, :] = list(map(float, f.readline().strip().replace(",", ".").split(" ")))

print("Processing..")
# samples to mV
mep_signal *= resolution
mep_signal *= (resolution / 1000)
# samples to ms
mep_time = np.arange(-1754, samples_count[0] - 1754) * sampling_interval / 1000

# find peaks for signals with N/A markers
markers_neg[markers_neg == "N/A"] = '0.0'
markers_pos[markers_pos == "N/A"] = '0.0'
markers_neg = np.array(list(map(float, markers_neg)))
markers_pos = np.array(list(map(float, markers_pos)))
markers_nans = (markers_neg == 0.0) | (markers_pos == 0.0)

# averaged signal
mep_mean = np.mean(mep_signal, axis=1)
markers_pos_mean = np.max(mep_mean[2000:])
markers_neg_mean = np.min(mep_mean[2000:])

interpeak_latency = np.zeros(mep_signal.shape[1])
p2p_amplitude = np.zeros(mep_signal.shape[1])
for idx in range(mep_signal.shape[1]):
    if markers_nans[idx]:
        interpeak_latency[idx] = np.nan
        p2p_amplitude[idx] = np.nan
    else:
        if markers_neg[idx] > markers_pos[idx]:
            interpeak_latency[idx] = markers_neg[idx] - markers_pos[idx]
        elif markers_neg[idx] < markers_pos[idx]:
            interpeak_latency[idx] = markers_pos[idx] - markers_neg[idx]
        else:
            interpeak_latency[idx] = 0
        p2p_amplitude[idx] = mep_signal[(np.where(mep_time == markers_pos[idx])[0][0]), idx] + np.abs(mep_signal[(np.where(mep_time == markers_neg[idx])[0][0]), idx])

# generate report
txt_filename = file_name.replace(".ascii", "-report.txt")
print(f"Saving TXT: {txt_filename}")
with open(txt_filename, "w") as f:
    print(f"Stimulator: {sw_name}", file=f)
    print(f"Version: {sw_version}", file=f)
    print(f"Subject: {subject}", file=f)
    print(f"Subject ID: {subject_id}", file=f)
    print(f"Record ID: {record_id}", file=f)
    print(f"Date: {record_created}", file=f)
    print(f"Method: {method}", file=f)
    print(f"Coil type: {coil_type[0]}\n", file=f)
    print(f"Markers latency unit: {marker_latency_unit}", file=f)
    print(f"Sampling interval: {sampling_interval}", file=f)
    print(f"Sampling interval unit: {sampling_interval_unit}", file=f)
    print(f"Signal resolution: {resolution}", file=f)
    print(f"Signal resolution unit: {resolution_unit}", file=f)
    print(f"Number of signals: {signal_count}", file=f)
    print(f"Number of signals with undetected markers: {np.sum(markers_nans)}", file=f)
    print(f"Number of samples per signal: {samples_count[0]}\n", file=f)
    for idx in range(mep_signal.shape[1]):
        print(f"Signal: {(idx + 1)}", file=f)
        print(f"\tStimulation intensity: {stim_intens[idx]}", file=f)
        if not markers_nans[idx]:
            print(f"\tA- latency: {markers_neg[idx]} ms", file=f)
            print(f"\tA- amplitude: {np.round(mep_signal[(np.where(mep_time == markers_neg[idx])[0][0]), idx], 4)} V", file=f)
            print(f"\tA+ latency: {markers_pos[idx]} ms", file=f)
            print(f"\tA+ amplitude: {np.round(mep_signal[(np.where(mep_time == markers_pos[idx])[0][0]), idx], 4)} V", file=f)
            print(f"\tInterpeak latency: {np.round(interpeak_latency[idx], 4)} ms", file=f)
            print(f"\tPeak-to-peak amplitude: {np.round(p2p_amplitude[idx], 4)} V", file=f)
        else:
            print(f"\tA-: not detected", file=f)
            print(f"\tA+: not detected", file=f)

# generate plots
print("Generating plots..")
pdf_filename = file_name.replace('.ascii', '-butterfly_averaged.pdf')
plt.figure(figsize=(8, 4))
plt.subplot(211)
# butterfly plot
for idx in range(mep_signal.shape[1]):
    plt.plot(mep_time, mep_signal[:, idx], linewidth="0.5")
plt.xlim([-10, 50])
plt.ylim([-1, 1])
plt.xlabel('Time [ms]')
plt.ylabel('MEP signal [V]')
plt.title(f"Butterfly plot of {mep_signal.shape[1]} signals")
plt.subplot(212)
# averaged plot
plt.plot(mep_time, mep_mean, linewidth="0.5")
plt.plot(mep_time[np.abs(mep_mean - markers_neg_mean).argmin()], mep_mean[np.abs(mep_mean - markers_neg_mean).argmin()], 'x', label='A−')
plt.plot(mep_time[np.abs(mep_mean - markers_pos_mean).argmin()], mep_mean[np.abs(mep_mean - markers_pos_mean).argmin()], 'x', label='A+')
plt.xlim([-10, 50])
plt.ylim([-1, 1])
plt.xlabel('Time [ms]')
plt.ylabel('MEP signal [V]')
plt.legend()
plt.title("Averaged signal")
plt.subplots_adjust(hspace=1)
plt.savefig(pdf_filename, format="pdf", bbox_inches="tight")

# plot individual signals with markers
for idx in range(0, mep_signal.shape[1]):
    plt.figure(figsize=(8, 4))
    plt.plot(mep_time, mep_signal[:, idx], linewidth="0.5")
    if not markers_nans[idx]:
        plt.plot(mep_time[(np.where(mep_time == markers_neg[idx])[0][0])], mep_signal[(np.where(mep_time == markers_neg[idx])[0][0]), idx], 'x', label='A−')
        plt.plot(mep_time[(np.where(mep_time == markers_pos[idx])[0][0])], mep_signal[(np.where(mep_time == markers_pos[idx])[0][0]), idx], 'x', label='A+')
        plt.legend()
    plt.xlim([-10, 50])
    plt.ylim([-1, 1])
    plt.xlabel('Time [ms]')
    plt.ylabel('MEP signal [V]')
    plt.title(f"Signal: {(idx + 1)}")
    pdf_filename = file_name.replace('.ascii', "-" + str(idx+1) + ".pdf")
    plt.savefig(pdf_filename, format="pdf", bbox_inches="tight")
    plt.close()

csv_filename = file_name.replace('.ascii', '-signals.csv')
print(f"Saving CSV: {csv_filename}")
mep_signal_df = pd.DataFrame({'time': mep_time})
for idx in range(0, mep_signal.shape[1]):
    signal_name = "signal_" + str(idx+1)
    mep_signal_df[signal_name] = mep_signal[:, idx]
mep_signal_df.to_csv(csv_filename, index=False)

csv_filename = file_name.replace('.ascii', '-peaks.csv')
print(f"Saving CSV: {csv_filename}")
mep_peaks_df = pd.DataFrame({'signal': range(1, mep_signal.shape[1]+1)})
peaks = np.zeros((mep_signal.shape[1], 4))
for idx in range(0, mep_signal.shape[1]):
    if not markers_nans[idx]:
        peaks[idx] = [markers_pos[idx], mep_signal[(np.where(mep_time == markers_pos[idx])[0][0]), idx], markers_neg[idx], mep_signal[(np.where(mep_time == markers_neg[idx])[0][0]), idx]]
    else:
        peaks[idx] = [np.nan, np.nan, np.nan, np.nan]
mep_peaks_df["apos_latency"] = peaks[:, 0]
mep_peaks_df["apos_amplitude"] = peaks[:, 1]
mep_peaks_df["aneg_latency"] = peaks[:, 2]
mep_peaks_df["aneg_amplitude"] = peaks[:, 3]
mep_peaks_df.to_csv(csv_filename, index=False)

print("Analysis completed.")
